import cv2
import os
import numpy as np


class ImCutMerge(object):
    def __init__(self):
        self.__exception_process = None
        self.__clip_rows = 0
        self.__clip_cols = 0
        self.__img = None
        self.__result_clip_dic = {}
        self.__TMP_RESULT_DIR = "tmp_result"
        self.__TMP_INFO_FILE = "tmp_info.txt"

    def set_exception_process(self, exception_process):
        """
        设置一个错误处理函数
        :param exception_process: 这要是一个错误处理函数其参数为出错的文本信息
        """
        self.__exception_process = exception_process

    def __creat_and_check_dir(self, result):
        """
        若目录不存在则创建目录
        :param result: 要创建的目录的名字
        """
        if not os.path.exists(str(result)):
            cwd = os.getcwd()
            cwd += "\\"
            os.mkdir(cwd + result)

    def __get_clip_info(self, img_shape, clip_rows, clip_cols):
        """
        :param img_shape: 总图的基本信息
        :param clip_rows: 切片的行数
        :param clip_cols: 切片的列数
        :return: 切片宽度，切片高度，最后一个切片的宽度，最后一个切片的高度
        """
        high = img_shape[0]  # 图片总高度
        width = img_shape[1]  # 图片总宽度
        clip_h = high // clip_rows  # 得到切片高度
        clip_w = width // clip_cols  # 得到切片宽度
        last_clip_h = clip_h - int(high % clip_h)
        last_clip_w = clip_w - int(width % clip_w)
        return clip_w, clip_h, last_clip_w, last_clip_h

    def img_cut(self, img, clip_rows, clip_cols, clip_process_callback=None):
        """

        :param img: 要分割的原始图像
        :param clip_rows: 分割的行数
        :param clip_cols: 分割的列数
        :param clip_process_callback: 分割后单张小图片的处理回调 其参数为单张分割后的小图
        :return: 返回处理后的且合并后的大图
        """
        clip_width, clip_high, clip_last_width, clip_last_high = 0, 0, 0, 0
        if img is not None:
            self.__img = img
            clip_width, clip_high, clip_last_width, clip_last_high = \
                self.__get_clip_info(img.shape, clip_rows, clip_cols)
            if clip_width <= 1 or clip_high <= 1:
                self.__exception_process("切片数量超限")
            # clip_top, clip_bottom, clip_left, clip_right = 0, clip_high, 0, clip_width
            self.__clip_rows = clip_rows
            self.__clip_cols = clip_cols
            for i in range(0, clip_rows):
                for j in range(0, clip_cols):
                    clip_top = clip_high * i
                    # if判断最后一片做特殊处理
                    if i == (clip_rows - 1):
                        clip_bottom = clip_top + clip_last_high
                    else:
                        clip_bottom = clip_top + clip_high
                    clip_left = clip_width * j
                    # 同上
                    if j == (clip_cols - 1):
                        clip_right = clip_left + clip_last_width
                    else:
                        clip_right = clip_left + clip_width
                    temp_img = img[clip_top:clip_bottom, clip_left:clip_right]
                    if clip_process_callback is not None:
                        temp_img = clip_process_callback(temp_img)
                        self.__result_clip_dic[(i, j)] = temp_img
                    else:
                        # 默认方式，回调函数为空时
                        self.__creat_and_check_dir(self.__TMP_RESULT_DIR)
                        clip_name = "%d_%d.jpg" % (i, j)
                        cv2.imwrite("./%s\\%s" % (self.__TMP_RESULT_DIR, clip_name), temp_img)
        else:
            self.__exception_process("img为空")

        if clip_process_callback is not None:
            self.__img_merge_list()
            return self.__img
        else:
            file = open("./%s\\%s" % (self.__TMP_RESULT_DIR, self.__TMP_INFO_FILE), 'w')
            file.write("%d\n" % self.__img.shape[0])
            file.write("%d\n" % self.__img.shape[1])
            file.write("%d\n" % self.__img.shape[2])
            file.write("%d\n" % clip_cols)
            file.write("%d\n" % clip_rows)
            file.write("%d\n" % clip_width)
            file.write("%d\n" % clip_high)
            file.write("%d\n" % clip_last_width)
            file.write("%d\n" % clip_last_high)
            file.flush()
            file.close()
            del img
            return None

    def img_merge_dir(self):
        filepath = "./%s\\" % self.__TMP_RESULT_DIR
        if os.path.exists(filepath + self.__TMP_INFO_FILE):
            file = open(filepath + self.__TMP_INFO_FILE, 'r')
            tl = []
            for i in file:
                tl.append(int(i))
            mhigt, mwidth, mch, mclip_cols, mclip_rows, \
            mclip_width, mclip_high, mclip_last_width, mclip_last_high = tl
            result_img = np.zeros((mhigt, mwidth, mch), dtype=np.uint8)
            # 注意设置位深度
            temp_file_list = []
            for i in os.listdir(self.__TMP_RESULT_DIR):
                if i.endswith(".jpg"):
                    temp_file_list.append(i)
            # print(len(temp_file_list))
            if len(temp_file_list) < mclip_rows * mclip_cols:
                self.__exception_process("图片数量不足")
            else:

                for tf in temp_file_list:
                    i, j = int(tf.split('_')[0]), int(tf.split('_')[1].split('.')[0])
                    clip_top = mclip_high * i
                    # if判断最后一片做特殊处理
                    if i == (mclip_rows - 1):
                        clip_bottom = clip_top + mclip_last_high
                    else:
                        clip_bottom = clip_top + mclip_high
                    clip_left = mclip_width * j
                    # 同上
                    if j == (mclip_cols - 1):
                        clip_right = clip_left + mclip_last_width
                    else:
                        clip_right = clip_left + mclip_width
                    tm_img = cv2.imread(filepath + tf)
                    # cv2.imshow("test", tm_img)
                    # print(result_img.shape)
                    result_img[clip_top:clip_bottom, clip_left:clip_right] = tm_img
                return result_img
        else:
            self.__exception_process("信息文件不存在")

    def __img_merge_list(self):
        """
        完成图像合并，使用字典的方式暂存分割后的切片图像，此函数从字典中取得数据
        """
        if len(self.__result_clip_dic) != 0:
            clip_width, clip_high, clip_last_width, clip_last_high = \
                self.__get_clip_info(self.__img.shape, self.__clip_rows, self.__clip_cols)
            n = 0
            for i in range(0, self.__clip_rows):
                for j in range(0, self.__clip_cols):
                    clip_top = clip_high * i
                    # if判断最后一片做特殊处理
                    if i == (self.__clip_rows - 1):
                        clip_bottom = clip_top + clip_last_high
                    else:
                        clip_bottom = clip_top + clip_high
                    clip_left = clip_width * j
                    # 同上
                    if j == (self.__clip_cols - 1):
                        clip_right = clip_left + clip_last_width
                    else:
                        clip_right = clip_left + clip_width
                    self.__img[clip_top:clip_bottom, clip_left:clip_right] = self.__result_clip_dic[(i, j)]
                    # cv2.imshow("test",self.__result_clip_list.pop(n))
                    n = n + 1
        else:
            self.__exception_process("结果列表为空")


im_cut_merge = ImCutMerge()
